
import math
import torch
import torch.nn.functional as F

def pixel_centers(H, W, device):
    ys = torch.arange(H, device=device) + 0.5
    xs = torch.arange(W, device=device) + 0.5
    yy, xx = torch.meshgrid(ys, xs, indexing="ij")
    return torch.stack([xx, yy], dim=-1)  # (H,W,2)

def downsample_bicubic(img, scale):
    """
    Downsample an image tensor by the given scale using bicubic interpolation.

    Args:
        img (torch.Tensor): Image tensor of shape (C, H, W) with values in [0, 1].
        scale (float): Downsampling factor (> 1.0), resulting size will be approximately (H/scale, W/scale).

    Returns:
        torch.Tensor: Downsampled image tensor of shape (C, h, w).
    """
    C, H, W = img.shape
    h = max(1, int(round(H/scale)))
    w = max(1, int(round(W/scale)))
    img = img.unsqueeze(0)
    lr = F.interpolate(img, size=(h,w), mode="bicubic", align_corners=False)
    return lr[0]

def sample_hr_pixels(Hh, Wh, N, device):
    ys = torch.randint(0, Hh, (N,), device=device)
    xs = torch.randint(0, Wh, (N,), device=device)
    return ys, xs  # integer pixel indices

def gather_pixels(img, ys, xs):
    """
    Gather pixel values at integer indices.

    Args:
        img (torch.Tensor): (B, C, H, W)
        ys (torch.Tensor): (B, N) or (N,)
        xs (torch.Tensor): (B, N) or (N,)

    Returns:
        torch.Tensor: (B, C, N) sampled pixels.
    """
    B, C, H, W = img.shape
    if ys.dim() == 1:
        ys = ys.unsqueeze(0).expand(B,-1)
        xs = xs.unsqueeze(0).expand(B,-1)
    idx = ys * W + xs
    flat = img.view(B, C, -1)
    out = torch.gather(flat, 2, idx.unsqueeze(1).expand(B,C,-1))
    return out

def psnr(pred, target, eps=1e-8):
    mse = torch.mean((pred - target) ** 2)
    return -10.0 * torch.log10(mse + eps)
